import wenet
import torch
import librosa
import numpy as np
import torchaudio
import os

# 检查是否有GPU可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 加载模型
model = wenet.load_model('chinese')

# 确保模型内部所有组件都在正确的设备上
# 注意：wenet模型可能有其自身的设备管理方式，我们不直接修改它

# 使用librosa加载音频文件（可处理PCM_24格式）
def load_audio_file(file_path, sample_rate=16000):
    """使用librosa加载音频文件，并转换为torchaudio所需格式"""
    try:
        # 使用librosa加载音频，会自动处理PCM_24格式
        waveform, sr = librosa.load(file_path, sr=sample_rate, mono=True)
        
        # 转换为torch张量并调整维度
        waveform_tensor = torch.FloatTensor(waveform).unsqueeze(0)
        
        print(f"音频已加载，长度：{len(waveform)}，采样率：{sr}Hz")
        return waveform_tensor, sr
    except Exception as e:
        print(f"加载音频文件时出错: {e}")
        return None, None

# 确保temp目录存在
os.makedirs('input_files/temp', exist_ok=True)

# 音频路径
audio_path = 'input_files/sichuan_dialet/000001.wav'

# 使用自定义加载方法处理音频
print("使用自定义方法加载音频...")
waveform, sample_rate = load_audio_file(audio_path)

if waveform is not None:
    # 确保音频格式正确
    if waveform.dim() == 2 and waveform.size(0) == 1:
        print("准备转写音频...")
        
        # 将音频保存为新的临时文件（PCM_16格式）
        temp_audio_path = 'input_files/temp/temp_audio.wav'
        torchaudio.save(temp_audio_path, waveform, sample_rate, format="wav", 
                       encoding="PCM_S", bits_per_sample=16)
        print(f"临时文件已保存: {temp_audio_path}")
        
        # 直接使用模型的compute_feats方法处理音频特征
        try:
            # 查看模型内部源码，开始手动执行转写流程
            print("正在提取音频特征...")
            
            # 检查模型是否有compute_feats方法
            if hasattr(model, 'compute_feats'):
                # 使用修改后的方法直接进行识别
                print("执行识别...")
                result = model.transcribe(temp_audio_path)
                print(f"识别结果: {result['text']}")
            else:
                print("模型没有compute_feats方法，无法继续处理")
        except Exception as e:
            print(f"处理音频时发生错误: {e}")
            
            # 尝试修改wenet源码或者使用另一种方法
            print("请尝试如下操作：")
            print("1. 使用ffmpeg转换音频文件为16bit PCM格式")
            print("    ffmpeg -i input_files/sichuan_dialet/000001.wav -acodec pcm_s16le input_files/temp/audio16bit.wav")
            print("2. 直接使用转换后的音频文件: model.transcribe('input_files/temp/audio16bit.wav')")
    else:
        print("音频格式不正确")
else:
    print("无法加载音频文件")